package cn.feizhu.aimarket.ai.workflow.node;

import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import cn.feizhu.aimarket.ai.workflow.model.enums.NodeStatus;
import cn.feizhu.aimarket.model.dto.ai.outline.NovelOutlineResponse;
import cn.hutool.json.JSONUtil;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/**
 * 聚合节点
 * 将多个分布式文档合并成完整的内容
 *
 * @author feizhu
 * @since 2025/10/16
 */
public class AggregateNode implements NodeAction {

    private static final Logger logger = LoggerFactory.getLogger(AggregateNode.class);
    
    private final Map<String, NodeStatus> node2Status;
    
    public static final String NODE_NAME = "aggregate";

    public AggregateNode(Map<String, NodeStatus> node2Status) {
        this.node2Status = node2Status;
    }

    @Override
    public Map<String, Object> apply(OverAllState state) {
        logger.info("=== 聚合节点开始执行 ===");
        node2Status.put(NODE_NAME, NodeStatus.RUNNING);

        try {
            // 获取生成的多个文档
            @SuppressWarnings("unchecked")
            List<String> subDocs = (List<String>) state.value(WorkflowControlConstant.SUB_DOCS)
                    .orElse(List.of());
            // TODO 注意这里取出来是一个LinkedHashMap，需要两次转换
            Object outlineResponseRaw = state.value(WorkflowControlConstant.OUTLINE).orElse(new NovelOutlineResponse());
            NovelOutlineResponse outlineResponse;
            if(outlineResponseRaw instanceof LinkedHashMap<?,?>) {
                String jsonStr = JSONUtil.toJsonStr(outlineResponseRaw);
                outlineResponse = JSONUtil.toBean(jsonStr, NovelOutlineResponse.class);
            } else {
                outlineResponse = (NovelOutlineResponse) outlineResponseRaw;
            }

            logger.info("开始聚合{}个文档部分", subDocs.size());

            // 构建完整的聚合文档 运行期最大长度为21亿
            StringBuilder aggregatedContent = new StringBuilder();
            for (String subDoc : subDocs) {
                aggregatedContent.append(subDoc).append("\n\n");
            }
            // TODO 其实这个节点没啥必要
            String finalContent = aggregatedContent.toString();

            node2Status.put(NODE_NAME, NodeStatus.COMPLETED);
            return Map.of(
                    WorkflowControlConstant.AGGREGATED_CONTENT, finalContent,
                    WorkflowControlConstant.AGGREGATE_STATS, Map.of(
                            "sub_doc_count", subDocs.size(),
                            "outline_theme", outlineResponse.getOverview() != null ? outlineResponse.getOverview().getTheme() : "N/A",
                            "aggregation_time", LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME)
                    )
            );

        } catch (Exception e) {
            logger.error("聚合节点执行异常", e);
            node2Status.put(NODE_NAME, NodeStatus.FAILED);
            return Map.of(
                    "aggregated_content", "文档聚合失败：" + e.getMessage(),
                    "aggregate_stats", Map.of("error", e.getMessage())
            );
        } finally {
            logger.info("=== 聚合节点执行完成 ===");
        }
    }
}
